#!/usr/bin/env python3
"""
示例数据下载脚本

功能：
- 下载公开的车牌识别测试数据
- 提供数据集获取指南
- 生成合成车牌数据
- 验证下载的数据

作者：车牌识别系统
日期：2024年
"""

import os
import sys
import requests
import yaml
from pathlib import Path
from urllib.parse import urlparse
import logging
from typing import Dict, List, Optional
import json
import hashlib

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.utils.logger import Logger
from src.utils.config import ConfigManager

class DatasetDownloader:
    """数据集下载器类"""
    
    def __init__(self, config_path: str = "config.yaml"):
        """
        初始化下载器
        
        参数:
            config_path: 配置文件路径
        """
        self.config = ConfigManager(config_path)
        self.logger = Logger().get_logger("DatasetDownloader")
        
        # 创建数据目录
        self.data_dir = Path("data")
        self.images_dir = self.data_dir / "images"
        self.videos_dir = self.data_dir / "videos"
        
        for dir_path in [self.data_dir, self.images_dir, self.videos_dir]:
            dir_path.mkdir(exist_ok=True)
            
    def download_file(self, url: str, filename: str, description: str = "") -> bool:
        """
        下载单个文件
        
        参数:
            url: 下载链接
            filename: 保存文件名
            description: 文件描述
            
        返回:
            bool: 下载是否成功
        """
        try:
            self.logger.info(f"开始下载: {description}")
            self.logger.info(f"URL: {url}")
            
            response = requests.get(url, stream=True, timeout=30)
            response.raise_for_status()
            
            file_path = self.data_dir / filename
            with open(file_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        
            self.logger.info(f"下载完成: {file_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"下载失败 {url}: {e}")
            return False
            
    def download_test_images(self) -> None:
        """下载测试图片"""
        self.logger.info("开始下载测试图片...")
        
        sample_data = self.config.get('sample_data', {})
        test_images = sample_data.get('images', {}).get('test_images', [])
        
        success_count = 0
        total_count = len(test_images)
        
        for i, image_info in enumerate(test_images):
            url = image_info.get('url', '')
            description = image_info.get('description', f'测试图片{i+1}')
            
            if not url:
                continue
                
            # 从URL解析文件名
            parsed_url = urlparse(url)
            filename = f"test_image_{i+1}.jpg"
            
            if self.download_file(url, f"images/{filename}", description):
                success_count += 1
                
        self.logger.info(f"测试图片下载完成: {success_count}/{total_count}")
        
    def show_dataset_guide(self) -> None:
        """显示数据集获取指南"""
        print("\n" + "="*60)
        print("🚗 车牌识别数据集获取指南")
        print("="*60)
        
        sample_data = self.config.get('sample_data', {})
        guide = sample_data.get('acquisition_guide', {})
        
        # 学术数据集
        academic_datasets = guide.get('academic_datasets', [])
        if academic_datasets:
            print("\n📚 学术数据集:")
            for dataset in academic_datasets:
                print(f"\n• {dataset.get('name', 'Unknown')}")
                print(f"  描述: {dataset.get('description', 'N/A')}")
                if 'contact' in dataset:
                    print(f"  联系方式: {dataset['contact']}")
                if 'requirement' in dataset:
                    print(f"  申请要求: {dataset['requirement']}")
                if 'paper' in dataset:
                    print(f"  论文链接: {dataset['paper']}")
                    
        # 可直接下载的数据集
        images = sample_data.get('images', {})
        print("\n💾 可直接下载的数据集:")
        
        for key, dataset in images.items():
            if key == 'test_images':
                continue
            print(f"\n• {dataset.get('description', key)}")
            print(f"  链接: {dataset.get('url', 'N/A')}")
            print(f"  许可证: {dataset.get('license', 'N/A')}")
            if 'count' in dataset:
                print(f"  数量: {dataset['count']} 张")
            if 'note' in dataset:
                print(f"  注意: {dataset['note']}")
                
        # 自建数据集建议
        diy_tips = guide.get('diy_dataset', {}).get('tips', [])
        if diy_tips:
            print("\n🛠️ 自建数据集建议:")
            for tip in diy_tips:
                print(f"  • {tip}")
                
        # 开源工具
        tools = guide.get('tools', [])
        if tools:
            print("\n🔧 推荐开源工具:")
            for tool in tools:
                print(f"\n• {tool.get('name', 'Unknown')}")
                print(f"  链接: {tool.get('url', 'N/A')}")
                print(f"  描述: {tool.get('description', 'N/A')}")
                
    def generate_synthetic_data(self, count: int = 10) -> None:
        """
        生成合成车牌数据
        
        参数:
            count: 生成数量
        """
        try:
            from PIL import Image, ImageDraw, ImageFont
            import random
            import string
            
            self.logger.info(f"开始生成 {count} 张合成车牌图片...")
            
            # 中国车牌省份简称
            provinces = ['京', '津', '沪', '渝', '冀', '豫', '云', '辽', '黑', '湘', 
                        '皖', '鲁', '新', '苏', '浙', '赣', '鄂', '桂', '甘', '晋', 
                        '蒙', '陕', '吉', '闽', '贵', '粤', '青', '藏', '川', '宁', '琼']
            
            for i in range(count):
                # 生成车牌号
                province = random.choice(provinces)
                letter = random.choice(string.ascii_uppercase)
                numbers = ''.join(random.choices(string.digits + string.ascii_uppercase, k=5))
                plate_number = f"{province}{letter}{numbers}"
                
                # 创建车牌图片
                img = Image.new('RGB', (240, 80), color='blue')
                draw = ImageDraw.Draw(img)
                
                try:
                    # 尝试使用系统字体
                    font = ImageFont.truetype("arial.ttf", 24)
                except:
                    # 使用默认字体
                    font = ImageFont.load_default()
                
                # 绘制白色边框
                draw.rectangle([5, 5, 235, 75], outline='white', width=2)
                
                # 绘制车牌号
                draw.text((20, 25), plate_number, fill='white', font=font)
                
                # 保存图片
                filename = f"synthetic_plate_{i+1:03d}_{plate_number}.png"
                img.save(self.images_dir / filename)
                
            self.logger.info(f"合成车牌生成完成: {count} 张")
            
        except ImportError:
            self.logger.warning("PIL库未安装，无法生成合成数据。请运行: pip install Pillow")
        except Exception as e:
            self.logger.error(f"生成合成数据失败: {e}")
            
    def create_dataset_info(self) -> None:
        """创建数据集信息文件"""
        info = {
            "dataset_name": "车牌识别测试数据集",
            "created_date": str(Path().cwd()),
            "description": "用于车牌识别系统测试的示例数据",
            "structure": {
                "images/": "测试图片目录",
                "videos/": "测试视频目录",
                "synthetic/": "合成数据目录"
            },
            "sources": {
                "OpenALPR": "https://github.com/openalpr/openalpr",
                "HuggingFace": "https://huggingface.co/datasets/UniDataPro/license-plate-detection",
                "UFPR-ALPR": "https://web.inf.ufpr.br/vri/databases/ufpr-alpr/"
            },
            "usage": {
                "测试识别算法": "使用images目录下的图片测试车牌识别准确率",
                "训练模型": "可以使用这些数据作为训练集的补充",
                "算法验证": "验证不同场景下的识别效果"
            },
            "license": "仅用于学术研究和个人学习",
            "contact": "如有问题请查看项目文档"
        }
        
        info_file = self.data_dir / "dataset_info.json"
        with open(info_file, 'w', encoding='utf-8') as f:
            json.dump(info, f, ensure_ascii=False, indent=2)
            
        self.logger.info(f"数据集信息已保存: {info_file}")
        
    def validate_downloads(self) -> None:
        """验证下载的数据"""
        self.logger.info("验证下载的数据...")
        
        image_count = len(list(self.images_dir.glob("*.jpg"))) + len(list(self.images_dir.glob("*.png")))
        video_count = len(list(self.videos_dir.glob("*.mp4")))
        
        print(f"\n📊 数据统计:")
        print(f"  图片数量: {image_count}")
        print(f"  视频数量: {video_count}")
        print(f"  数据目录: {self.data_dir.absolute()}")
        
        if image_count > 0:
            print(f"\n✅ 成功下载 {image_count} 张测试图片")
        else:
            print(f"\n⚠️  未找到测试图片，建议手动下载或生成合成数据")
            
def main():
    """主函数"""
    print("🚗 车牌识别系统 - 示例数据下载工具")
    print("="*50)
    
    try:
        downloader = DatasetDownloader()
        
        # 显示数据集指南
        downloader.show_dataset_guide()
        
        # 询问用户操作
        print("\n请选择操作:")
        print("1. 下载测试图片")
        print("2. 生成合成车牌数据")
        print("3. 创建数据集信息文件")
        print("4. 全部执行")
        print("5. 退出")
        
        choice = input("\n请输入选择 (1-5): ").strip()
        
        if choice == "1":
            downloader.download_test_images()
        elif choice == "2":
            count = input("请输入生成数量 (默认10): ").strip()
            count = int(count) if count.isdigit() else 10
            downloader.generate_synthetic_data(count)
        elif choice == "3":
            downloader.create_dataset_info()
        elif choice == "4":
            downloader.download_test_images()
            downloader.generate_synthetic_data(10)
            downloader.create_dataset_info()
        elif choice == "5":
            print("退出程序")
            return
        else:
            print("无效选择")
            return
            
        # 验证结果
        downloader.validate_downloads()
        
        print("\n✅ 操作完成！")
        print(f"数据已保存到: {downloader.data_dir.absolute()}")
        print("\n💡 提示:")
        print("- 可以将自己的测试图片放入 data/images/ 目录")
        print("- 可以将测试视频放入 data/videos/ 目录")
        print("- 查看 dataset_info.json 了解数据集详情")
        
    except KeyboardInterrupt:
        print("\n\n用户取消操作")
    except Exception as e:
        print(f"\n❌ 错误: {e}")
        logging.exception("程序执行出错")

if __name__ == "__main__":
    main() 